import pandas as pd
import paramiko
import time
import os
import re
import traceback
from openpyxl import load_workbook

def read_server_info(excel_path):
    """
    从Excel文件中读取服务器连接信息
    :param excel_path: Excel文件路径
    :return: 包含服务器信息的DataFrame
    """
    try:
        df = pd.read_excel(excel_path)
        required_columns = ['IP地址', '端口', '账户', '密码', 'root账户', 'root密码']

        # 检查必要的列是否存在
        for col in required_columns:
            if col not in df.columns:
                print(f"错误: Excel文件中缺少必要的列 '{col}'")
                return None

        # 添加备注列（如果不存在）
        if '备注' not in df.columns:
            df['备注'] = ''

        return df
    except Exception as e:
        print(f"读取Excel文件时出错: {e}")
        return None

def connect_to_server(ip, port, username, password):
    """
    连接到服务器
    :param ip: 服务器IP地址
    :param port: 端口
    :param username: 用户名
    :param password: 密码
    :return: SSH客户端连接和错误信息（如果有）
    """
    try:
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh.connect(hostname=ip, port=port, username=username, password=password, timeout=10)
        return ssh, None
    except Exception as e:
        error_msg = f"连接服务器失败: {str(e)}"
        print(f"连接服务器 {ip}:{port} 失败: {e}")
        return None, error_msg

def switch_to_root(ssh, root_username, root_password):
    """
    切换到root用户
    :param ssh: SSH连接
    :param root_username: root用户名
    :param root_password: root密码
    :return: 交互式shell通道和错误信息（如果有）
    """
    try:
        # 创建交互式shell
        channel = ssh.invoke_shell()
        time.sleep(1)

        # 执行su命令切换到root
        channel.send(f"su - {root_username}\n")
        time.sleep(1)

        # 发送root密码
        channel.send(f"{root_password}\n")
        time.sleep(2)

        # 检查是否成功切换到root
        channel.send("whoami\n")
        time.sleep(1)
        output = channel.recv(1024).decode('utf-8', errors='ignore')

        if root_username in output:
            print(f"成功切换到{root_username}用户")
            return channel, None
        else:
            error_msg = f"切换到{root_username}用户失败，可能是密码错误或权限不足"
            print(f"切换到{root_username}用户失败")
            return None, error_msg
    except Exception as e:
        error_msg = f"切换用户过程中出错: {str(e)}"
        print(f"切换到{root_username}用户时出错: {e}")
        return None, error_msg

def get_java_processes(channel):
    """
    获取Java进程信息
    :param channel: 交互式shell通道
    :return: Java进程信息字符串和错误信息（如果有）
    """
    try:
        # 清空缓冲区
        if channel.recv_ready():
            channel.recv(1024)

        # 执行ps命令查找Java进程
        channel.send("ps -ef | grep java | grep -v grep\n")
        time.sleep(2)

        # 读取输出
        output = ""
        while channel.recv_ready():
            output += channel.recv(1024).decode('utf-8', errors='ignore')

        # 解析输出，提取进程号、包名和路径
        java_processes = []
        for line in output.splitlines():
            if "java" in line and not line.strip().startswith("grep"):
                # 尝试提取PID
                parts = line.split()
                if len(parts) > 1:
                    pid = parts[1]

                    # 尝试提取jar包名或类名
                    jar_match = re.search(r'-jar\s+(\S+\.jar)', line)
                    class_match = re.search(r'(\w+\.)+\w+', line)

                    if jar_match:
                        package = jar_match.group(1)
                    elif class_match:
                        package = class_match.group(0)
                    else:
                        package = "未知"

                    # 尝试提取路径
                    path_match = re.search(r'(/\w+)+/\w+\.jar', line)
                    if path_match:
                        path = path_match.group(0)
                    else:
                        path = "未知"

                    java_processes.append(f"PID: {pid}, 包名: {package}, 路径: {path}")

        result = "\n".join(java_processes) if java_processes else "未找到Java进程"
        return result, None
    except Exception as e:
        error_msg = f"获取Java进程信息时出错: {str(e)}"
        print(f"获取Java进程信息时出错: {e}")
        return None, error_msg

def write_results_to_excel(excel_path, df):
    """
    将结果写入Excel文件
    :param excel_path: Excel文件路径
    :param df: 包含结果的DataFrame
    """
    try:
        # 创建一个新的Excel文件路径，避免覆盖原始文件
        result_path = os.path.splitext(excel_path)[0] + "_results" + os.path.splitext(excel_path)[1]
        
        # 方法1：使用新版pandas的ExcelWriter方式
        try:
            df.to_excel(result_path, index=False)
            print(f"结果已成功写入到 {result_path}")
            return
        except Exception as e1:
            print(f"使用方法1写入Excel文件时出错: {e1}")
        
        # 方法2：使用openpyxl直接操作
        try:
            # 创建一个新的工作簿
            from openpyxl import Workbook
            wb = Workbook()
            ws = wb.active
            
            # 添加标题行
            for col_idx, column_name in enumerate(df.columns, 1):
                ws.cell(row=1, column=col_idx, value=column_name)
            
            # 添加数据行
            for row_idx, row in enumerate(df.values, 2):
                for col_idx, value in enumerate(row, 1):
                    ws.cell(row=row_idx, column=col_idx, value=value)
            
            # 保存工作簿
            wb.save(result_path)
            print(f"使用openpyxl直接写入到 {result_path}")
            return
        except Exception as e2:
            print(f"使用方法2写入Excel文件时出错: {e2}")
        
        print("所有写入方法都失败，无法保存结果")
    except Exception as e:
        print(f"写入Excel文件时出错: {e}")
        print("无法保存结果到Excel文件。请检查文件是否被其他程序占用。")

def main():
    # 设置Excel文件路径
    #excel_path = input("请输入Excel文件路径: ")
    excel_path="F:\\服务器清单2.xlsx"
    if not os.path.exists(excel_path):
        print(f"错误: 文件 '{excel_path}' 不存在")
        return

    # 读取服务器信息
    server_df = read_server_info(excel_path)
    if server_df is None:
        return

    # 遍历每个服务器
    for index, row in server_df.iterrows():
        ip = row['IP地址']
        port = int(row['端口']) if not pd.isna(row['端口']) else 22
        username = row['账户']
        password = row['密码']
        root_username = row['root账户']
        root_password = row['root密码']

        print(f"\n处理服务器: {ip}:{port}")
        
        # 记录当前服务器的处理状态
        status = {
            "成功": False,
            "错误信息": "",
            "处理阶段": "",
            "Java进程": ""
        }

        # 检查普通账户信息是否完整
        if pd.isna(username) or pd.isna(password):
            if pd.isna(root_username) or pd.isna(root_password):
                error_msg = f"服务器 {ip} 的账户信息不完整，无法连接"
                print(f"错误: {error_msg}")
                server_df.at[index, '备注'] = error_msg
                continue
            else:
                # 直接使用root账户
                username = root_username
                password = root_password
                print(f"使用root账户直接连接: {username}")

        # 连接服务器
        status["处理阶段"] = "连接服务器"
        ssh, error = connect_to_server(ip, port, username, password)
        if ssh is None:
            status["错误信息"] = error
            error_detail = f"连接失败: {ip}:{port} 用户名: {username}\n错误详情: {error}"
            server_df.at[index, '备注'] = error_detail
            print(error_detail)
            continue

        # 如果使用的是普通账户且root账户信息完整，则切换到root
        if username != root_username and not pd.isna(root_username) and not pd.isna(root_password):
            status["处理阶段"] = "切换到root用户"
            channel, error = switch_to_root(ssh, root_username, root_password)
            if channel is None:
                ssh.close()
                status["错误信息"] = error
                error_detail = f"连接成功但切换到root用户失败: {ip}\n错误详情: {error}"
                server_df.at[index, '备注'] = error_detail
                print(error_detail)
                continue
        else:
            # 直接使用当前连接创建通道
            try:
                channel = ssh.invoke_shell()
                time.sleep(1)
            except Exception as e:
                ssh.close()
                error_detail = f"创建交互式Shell失败: {ip}\n错误详情: {str(e)}"
                server_df.at[index, '备注'] = error_detail
                print(error_detail)
                continue

        # 获取Java进程信息
        status["处理阶段"] = "获取Java进程信息"
        java_info, error = get_java_processes(channel)
        
        if error:
            status["错误信息"] = error
            error_detail = f"获取Java进程信息失败: {ip}\n错误详情: {error}"
            server_df.at[index, '备注'] = error_detail
        else:
            status["成功"] = True
            status["Java进程"] = java_info
            server_df.at[index, '备注'] = java_info

        # 关闭连接
        try:
            channel.close()
            ssh.close()
        except:
            pass
        
        # 打印处理结果
        if status["成功"]:
            print(f"服务器 {ip} 处理完成，成功获取Java进程信息")
        else:
            print(f"服务器 {ip} 处理失败: {status['错误信息']}")

    # 将结果写入Excel
    write_results_to_excel(excel_path, server_df)
    print("所有服务器处理完成")

if __name__ == "__main__":
    main()